package cas.ihep.spark.rdd;

import com.google.common.base.Strings;
import org.apache.spark.*;
import org.apache.spark.api.java.JavaRDD;
import org.apache.spark.api.java.JavaSparkContext;
import org.apache.spark.rdd.RDD;
import scala.collection.Iterator;
import scala.collection.JavaConverters;
import scala.collection.Seq;
import scala.reflect.ClassManifestFactory;

import java.net.URI;
import java.util.*;

public class FileBasedRDD extends RDD<URI> {

    private static class FileBasedPartition implements Partition{
        private int _rddid;
        private int _index;
        private List<URI> uris;
        @Override
        public int index() {
            return _index;
        }

        public boolean equals(Object other){
            if(other instanceof FileBasedPartition){
                FileBasedPartition fbPartition=(FileBasedPartition)other;
                return _rddid==fbPartition._rddid && _index==fbPartition._index;
            }
            return false;
        }

    }

    public interface Spliter{
        String groupBy(URI uri);
    }

    private Partition _partitions[]=null;
    private List<URI> inputPaths;
    private int numOfSlices;
    private Spliter spliter;

    public static JavaRDD<URI> newRDD(List<URI> paths, Spliter s, JavaSparkContext jctx){
        return JavaRDD.fromRDD(new FileBasedRDD(paths,jctx.defaultParallelism(),s,jctx.sc(),null),ClassManifestFactory.fromClass(URI.class));
    }

    public static JavaRDD<URI> newRDD(List<URI> paths, int slices,Spliter s,JavaSparkContext jctx){
        return JavaRDD.fromRDD(new FileBasedRDD(paths,slices,s,jctx.sc(),null),ClassManifestFactory.fromClass(URI.class));
    }

    private FileBasedRDD(List<URI> filepaths,int slices,Spliter s, SparkContext sc, Seq<Dependency<?>> deps) {
        super(sc, deps, ClassManifestFactory.fromClass(URI.class));
        inputPaths=filepaths;
        numOfSlices=slices;
        spliter=s;
    }

    @Override
    public Iterator<URI> compute(Partition split, TaskContext context) {
        FileBasedPartition fbPartition=(FileBasedPartition)split;
        return new InterruptibleIterator<>(context, JavaConverters.asScalaIteratorConverter(fbPartition.uris.iterator()).asScala());
    }

    @Override
    public Partition[] getPartitions() {
        if(_partitions==null){
            HashMap<String,List<URI>> split=new HashMap<>();
            int sz=inputPaths.size();
            //noinspection ForLoopReplaceableByForEach
            for(int i=0;i<sz;i++){
                URI uri=inputPaths.get(i);
                String key= Strings.nullToEmpty(spliter.groupBy(uri));
                List<URI> uris=split.get(key);
                //noinspection Java8MapApi
                if(uris==null){
                    uris=new ArrayList<>();
                    split.put(key,uris);
                }
                uris.add(uri);
            }
            Set<Map.Entry<String,List<URI>>> splitset=split.entrySet();
            if(numOfSlices<1){
                numOfSlices=splitset.size();
            }
            int elementPerSlice=sz/numOfSlices;
            List<Partition> partitions=new ArrayList<>();
            int id=id();
            int index=0;
            for(Map.Entry<String,List<URI>> entry:splitset){
                List<List<URI>> lists=splitList(entry.getValue(),elementPerSlice);
                for(List<URI> part:lists){
                    FileBasedPartition fbPartition=new FileBasedPartition();
                    fbPartition._rddid=id;
                    fbPartition._index=index++;
                    fbPartition.uris=part;
                    partitions.add(fbPartition);
                }
            }
            _partitions= (Partition[]) partitions.toArray();
            if(_partitions.length!=numOfSlices){
                System.err.println("WARNING: The count of partitions("+_partitions.length+") is not equal to "+numOfSlices);
            }
        }
        return _partitions;
    }


    private static List<List<URI>> splitList(List<URI> resList,int count){
        if(resList==null ||count<1)
            return  null ;
        List<List<URI>> ret=new ArrayList<>();
        int size=resList.size();
        if(size<=count){
            ret.add(resList);
        }else{
            int pre=size/count;
            int last=size%count;
            for(int i=0;i<pre;i++){
                List<URI> itemList=new ArrayList<>();
                for(int j=0;j<count;j++){
                    itemList.add(resList.get(i*count+j));
                }
                ret.add(itemList);
            }
            if(last>0){
                List<URI> itemList=new ArrayList<>();
                for(int i=0;i<last;i++){
                    itemList.add(resList.get(pre*count+i));
                }
                ret.add(itemList);
            }
        }
        return ret;
    }
}
